{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epc\tpr\tre\tf1\tiou\tmiou\toa\tkappa\n",
      "\n",
      "Use load_from_local loader\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ad103b03f05646918c913ccaad3e7ffd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=1000.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "code_root = '/workspace/code/mmcd' # change it to /path/to/code_root\n",
    "data_root = '/workspace/dataset/stb' # change it to /path/to/data_root\n",
    "\n",
    "import sys,os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"7\"\n",
    "sys.path.insert(0,f'{code_root}/')\n",
    "\n",
    "import os\n",
    "import cv2\n",
    "\n",
    "import torch\n",
    "import torchvision.models as tvmodels\n",
    "import torch.nn.functional as F\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from torchvision.transforms import Resize\n",
    "from tqdm.notebook import tqdm\n",
    "from skimage import morphology\n",
    "\n",
    "\n",
    "from mmseg.datasets import TwoInputDataset,build_dataloader\n",
    "from mmseg.apis import init_segmentor, inference_segmentor, single_gpu_test\n",
    "from mmseg.utils import Metric\n",
    "\n",
    "\n",
    "def base_forward(model,x):\n",
    "#     x = model.extract_feat(x)\n",
    "#     y = model.decode_head(x)\n",
    "    y = model.encode_decode(x,{})\n",
    "    return y\n",
    "\n",
    "def forward(model,x,tta=False): # x=data['img']\n",
    "    if not tta:\n",
    "        return base_forward(model,x)\n",
    "    outs = []\n",
    "    origin_x = x.clone()\n",
    "    \n",
    "    y = base_forward(model,x)\n",
    "    \n",
    "    x_ = origin_x.flip(3)\n",
    "    y_ = base_forward(model,x_).flip(2)\n",
    "    y += F.softmax(y_,dim=1)\n",
    "    \n",
    "    x_ = origin_x.flip(4)\n",
    "    y_ = base_forward(model,x_).flip(3)\n",
    "    y += F.softmax(y_,dim=1)\n",
    "    \n",
    "    x_ = origin_x.transpose(3, 4).flip(4)\n",
    "    y_ = base_forward(model,x_).flip(3).transpose(2, 3)\n",
    "    y += F.softmax(y_,dim=1)\n",
    "    \n",
    "    x_ = origin_x.flip(4).transpose(3, 4)\n",
    "    y_ = base_forward(model,x_).transpose(2, 3).flip(3)\n",
    "    y += F.softmax(y_,dim=1)\n",
    "    \n",
    "    x_ = origin_x.flip(3).flip(4)\n",
    "    y_ = base_forward(model,x_).flip(3).flip(2)\n",
    "    y += F.softmax(y_,dim=1)\n",
    "    \n",
    "    return y/6.0\n",
    "\n",
    "flag_save = True # 保存结果\n",
    "flag_pre = True # resize\n",
    "flag_post = False\n",
    "flag_debug = False\n",
    "tta = True # 打开TTA\n",
    "t22t1 = False\n",
    "cal_metric = False\n",
    "save_submit = True # 保存提交格式的文件\n",
    "\n",
    "for isp in [576]:\n",
    "    \n",
    "    \n",
    "    name = 'upernet_hr40_256x256_stb'\n",
    "    name = 'upernet_hr40_512x512_stb'\n",
    "    name = 'ocrnet_hr40_512x512_stb'\n",
    "    name = 'upernet_hr40_576x576_stb'\n",
    "    # meta_file = f'{data_root}/stb/val.v1.txt'\n",
    "    meta_file = f'{code_root}/work_dirs/cd_stb/meta_files/test.txt'\n",
    "    config_file = f'{code_root}/work_dirs/cd_stb/{name}/{name}.py'\n",
    "    save_root = f'{code_root}/work_dirs/cd_stb/results_{name}_{isp}/'\n",
    "#     save_root = f'{code_root}/work_dirs/cd_stb/results_val_{name}_{isp}/'\n",
    "    checkpoint_root = f'{code_root}/work_dirs/cd_stb/{name}/'\n",
    "\n",
    "\n",
    "    f = open(f'{code_root}/work_dirs/cd_stb/result.txt','a')\n",
    "    f.write('='*88+'\\n')\n",
    "    f.write(f'{config_file}\\t{meta_file}\\t{checkpoint_root}\\n')\n",
    "    f.write('epc\\tpr\\tre\\tf1\\tiou\\tmiou\\toa\\tkappa\\n')\n",
    "    if flag_save:\n",
    "        os.system(f\"mkdir -p {save_root}\")\n",
    "    print('epc\\tpr\\tre\\tf1\\tiou\\tmiou\\toa\\tkappa\\n')\n",
    "#     for epc in range(160000,162000,2000):\n",
    "#     for epc in range(96000,122000,2000):\n",
    "#     for epc in [80000]:\n",
    "#     for epc in range(118000,162000,2000):\n",
    "#     for epc in range(5000,202000,5000):\n",
    "    for epc in [180000]:\n",
    "    \n",
    "        checkpoint_file = os.path.join(checkpoint_root,f'iter_{epc}.pth')\n",
    "        if not os.path.exists(checkpoint_file):\n",
    "            print(f'{checkpoint_file} not exists.')\n",
    "            continue\n",
    "\n",
    "        model = init_segmentor(config_file, checkpoint_file, device='cuda:0')\n",
    "        model = model.eval()\n",
    "\n",
    "        test_pipeline = [\n",
    "#             dict(type='Resize',height=256,width=256,p=1.0)\n",
    "            dict(type='Normalize',mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n",
    "            dict(type='ToTensorV2'),\n",
    "        ]\n",
    "\n",
    "        dataset = TwoInputDataset(meta_file=meta_file,sep='\\t',pipeline=test_pipeline,data_root=data_root)\n",
    "        dataloader = build_dataloader(dataset,1,1,shuffle=False)\n",
    "\n",
    "        metric = Metric(name='sn')\n",
    "        reisp = Resize([isp,isp])\n",
    "        \n",
    "        up512 = torch.nn.Upsample(size=(512,512), mode='bilinear',align_corners=False)\n",
    "        up513 = torch.nn.Upsample(size=(513,513), mode='bilinear',align_corners=False)\n",
    "        up256 = torch.nn.Upsample(size=(256,256), mode='bilinear',align_corners=False)\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            for idx,data in tqdm(enumerate(dataloader),total=len(dataset)):\n",
    "                for k,v in data.items():\n",
    "                    try:\n",
    "                        data[k] = v.cuda()\n",
    "                    except:\n",
    "                        pass\n",
    "                if t22t1:\n",
    "                    data['img'] = data['img'][:,[1,0],...]\n",
    "                if flag_pre:\n",
    "                    data['img'] = reisp(data['img'].squeeze()).unsqueeze(0)\n",
    "                    y = forward(model,data['img'],tta=tta)\n",
    "                    y = up512(y)\n",
    "                else:\n",
    "                    y = forward(model,data['img'],tta=tta)\n",
    "                    \n",
    "\n",
    "                if flag_post:\n",
    "                    out = y.argmax(dim=1)\n",
    "                    imo = out.clone().cpu().squeeze().numpy().astype(bool)\n",
    "                    imo = morphology.remove_small_holes(imo,area_threshold=1000,connectivity=8)\n",
    "                    imo = imo.astype(np.uint8)\n",
    "                    gtt = data['gt_semantic_seg'].cpu().squeeze().numpy()\n",
    "                    \n",
    "                else:\n",
    "                    out = y.argmax(dim=1)\n",
    "                    imo = out\n",
    "                    gtt = data['gt_semantic_seg']\n",
    "                \n",
    "                if cal_metric:\n",
    "                    metric(imo,gtt)\n",
    "                \n",
    "                if flag_save:\n",
    "                    line = dataset.data[idx]\n",
    "                    save_name = os.path.basename(line.strip().split('\\t')[-1]).split('.')[0]+'.png'\n",
    "                    save_path = os.path.join(save_root,save_name)\n",
    "                    factor = 1 if save_submit else 255\n",
    "                    cv2.imwrite(save_path,out.squeeze().detach().cpu().numpy()*factor)\n",
    "                if flag_debug:\n",
    "                    if idx==10:\n",
    "                        break\n",
    "\n",
    "\n",
    "        if cal_metric:\n",
    "            local=False\n",
    "            oa = metric.oa(local=local)\n",
    "            kappa = metric.kappa(local=local)\n",
    "            pr = metric.pr(local=local)\n",
    "            re = metric.re(local=local)\n",
    "            f1 = metric.f1(local=local)\n",
    "            iou = metric.iou(local=local)\n",
    "            miou = metric.miou(local=local)\n",
    "            print(f'{isp}\\t{epc}',pr.item(),re.item(),f1.item(),iou.item(),miou.item(),oa.item(),kappa.item(),sep='\\t')\n",
    "            f.write(f'{isp}\\t{epc}\\t{pr.item()}\\t{re.item()}\\t{f1.item()}\\t{iou.item()}\\t{miou.item()}\\t{oa.item()}\\t{kappa.item()}-----({idx})\\n')\n",
    "        f.flush()\n",
    "    f.write('='*88+'\\n')\n",
    "    f.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
